package cn.showcon.firstapp.service;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.OnnxEmbeddingModel;
import dev.langchain4j.model.embedding.onnx.PoolingMode;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

/**
 * @Author Xue Lanbin
 */
public class EmbeddingStoreExample {

    public static void main(String[] args) {
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
//        EmbeddingModel embeddingModel1 = new AllMiniLmL6V2EmbeddingModel();
        EmbeddingModel embeddingModel = createEmbeddingModel();

        TextSegment segment1 = TextSegment.from("我喜欢踢足球。");
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        embeddingStore.add(embedding1, segment1);

        TextSegment segment2 = TextSegment.from("今天的天气真好。");
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        embeddingStore.add(embedding2, segment2);

        Embedding queryEmbedding = embeddingModel.embed("你最喜爱的体育项目是什么?").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .minScore(0.0D)
                .build();

//        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(queryEmbedding, 1);
        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(embeddingSearchRequest).matches();
        EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);

        System.out.println(embeddingMatch.score());
        System.out.println(embeddingMatch.embedded().text());

        System.out.println(embeddingStore.serializeToJson());
    }

    public static EmbeddingModel createEmbeddingModel() {
        String pathToModel = "D:\\xuelb\\llm\\resources\\model\\embedding\\shibing624_text2vec-base-chinese\\onnx\\model.onnx";
        String pathToTokenizer = "D:\\xuelb\\llm\\resources\\model\\embedding\\shibing624_text2vec-base-chinese\\onnx\\tokenizer.json";
        PoolingMode poolingMode = PoolingMode.MEAN;
        OnnxEmbeddingModel localEmbeddingModel = new OnnxEmbeddingModel(pathToModel, pathToTokenizer, poolingMode);
        System.out.println("localEmbeddingModel.dimension(): " + localEmbeddingModel.dimension());
        return localEmbeddingModel;
    }
}
